import pathlib
import unittest

from src.processors import (
    ChineseEmotionSmallClassifier,
    Emotion,
    EmotionalClassificationProcessor,
    EmotionClassifier,
    SimpleEmotionClassifier,
)
from src.unmarshalers import User, UserStage

PROJECT_ROOT = pathlib.Path(__file__).parent.parent
CACHE_DIR = PROJECT_ROOT / "cache"


class FakeEmotionClassifier(EmotionClassifier):
    def classify(self, sentence: str) -> Emotion:
        if sentence == "bad":
            return Emotion.NEGATIVE
        return Emotion.POSITIVE


class EmotionClassifierTestCase(unittest.TestCase):
    def test_emotion_classification_processor(self):
        processor = EmotionalClassificationProcessor(FakeEmotionClassifier())
        result = processor.process(
            UserStage(
                [
                    User(
                        username="test",
                        posts=["good", "bad"],
                    )
                ]
            )
        )
        self.assertEqual(result.emotions[0][0], 0.5)
        self.assertEqual(result.emotions[0][1], 0)
        self.assertEqual(result.emotions[0][2], 0.5)

    def test_chinese_emotion_small_classifier(self):
        classifier = ChineseEmotionSmallClassifier(cache_dir=CACHE_DIR)
        self.assertEqual(classifier.classify("今天吃了顿大餐，好开心"), Emotion.POSITIVE)
        self.assertEqual(classifier.classify("今天被老板骂了，好难过"), Emotion.NEGATIVE)

    def test_simple_emotion_classifier(self):
        classifier = SimpleEmotionClassifier(cache_dir=CACHE_DIR)
        self.assertEqual(classifier.classify("今天吃了顿大餐，好开心"), Emotion.POSITIVE)
        self.assertEqual(classifier.classify("今天被老板骂了，好难过"), Emotion.NEGATIVE)
